#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
RL 数据集“正常样本”按比例移除脚本
================================

【功能说明】
- 读入一个 RL 阶段的 Parquet 数据文件（形如 med_pub_rl_train.parquet / med_pub_rl_test.parquet）。
- 识别其中的“正常眼底”样本，并按给定比例 r ∈ [0,1] 随机移除（r=1 表示全部移除）。
- 在输入文件同目录下，生成新数据文件：<原文件名>_rmNormal<r格式>.parquet
  例如：med_pub_rl_train_rmNormal0p8.parquet
- 同时输出一份统计 JSON：<原文件名>_rmNormal<r格式>_stats.json，记录移除前/后分布。

【正常样本判定规则（与现有 RL 生成逻辑兼容）】
- RL 生成时 ground_truth = "综上，本图片诊断为" + <标签串>；
- 我们解析 ground_truth 中 “诊断为” 之后的尾部，按照 [、 ， , 空白] 切分为标签列表；
- 若该标签集合恰好为 {"正常眼底"}，则视为“正常样本”。

【使用示例】
python remove_normal_from_parquet.py \
  --input_parquet ./experiments/dataset/RL_parquet/med_pub_rl_train.parquet \
  --remove_ratio 0.8 \
  --seed 42

【输出】
- ./experiments/dataset/RL_parquet/med_pub_rl_train_rmNormal0p8.parquet
- ./experiments/dataset/RL_parquet/med_pub_rl_train_rmNormal0p8_stats.json

作者：zym1105
时间：2025-08-15
"""

import argparse
import json
import os
import random
import re
from typing import List

import pandas as pd

# ===== 常量 =====
SYSTEM_ANSWER_PREFIX = "综上，本图片诊断为"
SEP_PATTERN = r"[,\s，、]+"  # 标签切分分隔符：英文逗号/中文逗号/顿号/空白


# ===== 工具函数 =====
def parse_args():
    parser = argparse.ArgumentParser(description="按比例移除 RL Parquet 数据中的正常样本")
    parser.add_argument(
        "--input_parquet",
        type=str,
        default="./experiments/dataset/RL_parquet/med_pub_rl_train.parquet",
        help="输入的 Parquet 文件路径",
    )
    parser.add_argument(
        "--remove_ratio",
        type=float,
        default=1.0,
        help="要移除的正常样本比例 r ∈ [0,1]，例如 0.8 表示移除 80% 正常样本；1.0 表示全部移除",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="随机种子（用于随机选择需要移除的正常样本）",
    )
    return parser.parse_args()


def safe_ratio_tag(r: float) -> str:
    """将浮点比例转换为文件名安全的后缀，例如 0.8 -> 0p8, 1.0 -> 1p0"""
    s = f"{r:.4g}"  # 简洁表示
    return s.replace(".", "p")


def extract_labels_from_gt(ground_truth: str) -> List[str]:
    """
    从 ground_truth 文本中提取标签列表。
    期望格式： "综上，本图片诊断为" + <标签串>，
    其中 <标签串> 可能是 "正常眼底" 或 "疾病1, 疾病2" 或 "疾病1、疾病2" 等。
    """
    if not isinstance(ground_truth, str):
        return []

    # 去掉前缀（若不存在前缀也尽量兼容）
    if "诊断为" in ground_truth:
        tail = ground_truth.split("诊断为", 1)[-1]
    else:
        # 兜底：直接用全文尝试切分
        tail = ground_truth

    tail = tail.strip()
    if not tail:
        return []

    # 切分为标签列表，去掉空字符串
    parts = re.split(SEP_PATTERN, tail)
    labels = [p.strip() for p in parts if p.strip()]
    return labels


def is_normal_sample(row) -> bool:
    """
    判定一条样本是否为“正常眼底”：
    - 优先从 row['reward_model'] 字段的 ground_truth 提取；
    - ground_truth 尾部按分隔符切分后，若集合恰好为 {"正常眼底"} 则认为是正常样本。
    - 若字段缺失或格式不符，返回 False。
    """
    rm = row.get("reward_model", None)
    if not isinstance(rm, dict):
        return False
    gt = rm.get("ground_truth", "")
    labels = extract_labels_from_gt(gt)
    return len(labels) == 1 and labels[0] == "正常眼底"


def main():
    args = parse_args()
    assert 0.0 <= args.remove_ratio <= 1.0, "remove_ratio 必须在 [0,1] 区间内"

    in_path = args.input_parquet
    assert os.path.exists(in_path), f"输入文件不存在：{in_path}"

    # 读取 parquet
    df = pd.read_parquet(in_path)
    total = len(df)
    print(f"[INFO] 加载样本：{total} 条  <- {in_path}")

    # 标记正常样本
    normal_mask = df.apply(is_normal_sample, axis=1)
    n_normal = int(normal_mask.sum())
    n_abnormal = total - n_normal
    print(f"[INFO] 初始分布：normal={n_normal}，abnormal={n_abnormal}，total={total}")

    # 计算需要移除的数量
    remove_ratio = args.remove_ratio
    n_remove = int(n_normal * remove_ratio + 1e-8)

    if n_remove == 0:
        print("[INFO] 移除数量为 0，直接复制输出。")
        filtered_df = df.copy()
        removed_idx = []
    else:
        # 从正常样本索引中随机选择要移除的样本
        random.seed(args.seed)
        normal_indices = df[normal_mask].index.tolist()
        removed_idx = random.sample(normal_indices, k=n_remove)
        filtered_df = df.drop(index=removed_idx)

    # 统计移除后的分布
    total_after = len(filtered_df)
    n_normal_after = int(filtered_df.apply(is_normal_sample, axis=1).sum())
    n_abnormal_after = total_after - n_normal_after

    print(f"[INFO] 计划移除正常样本：{n_remove} 条")
    print(f"[INFO] 结果分布：normal={n_normal_after}，abnormal={n_abnormal_after}，total={total_after}")

    # 组织输出路径
    base_dir = os.path.dirname(os.path.abspath(in_path))
    base_name = os.path.splitext(os.path.basename(in_path))[0]
    ratio_tag = safe_ratio_tag(remove_ratio)

    out_parquet = os.path.join(base_dir, f"{base_name}_rmNormal{ratio_tag}.parquet")
    out_stats = os.path.join(base_dir, f"{base_name}_rmNormal{ratio_tag}_stats.json")

    # 保存新 parquet
    filtered_df.to_parquet(out_parquet, index=False)
    print(f"[SUCCESS] 新数据已保存：{out_parquet}")

    # 保存统计 JSON（含移除样本数与索引）
    stats = {
        "input_parquet": in_path,
        "output_parquet": out_parquet,
        "remove_ratio": remove_ratio,
        "seed": args.seed,
        "before": {
            "total": total,
            "normal": n_normal,
            "abnormal": n_abnormal,
        },
        "removed": {
            "count": n_remove,
        },
        "after": {
            "total": total_after,
            "normal": n_normal_after,
            "abnormal": n_abnormal_after,
        },
    }
    with open(out_stats, "w", encoding="utf-8") as f:
        json.dump(stats, f, ensure_ascii=False, indent=2)
    print(f"[SUCCESS] 统计信息已保存：{out_stats}")


if __name__ == "__main__":
    main()
